"""
YASA睡眠分析器 - 优化版

本文件实现基于YASA库的脑波数据睡眠分析功能，包括：
- 使用YASA的SleepStaging类进行专业睡眠分期
- 智能30秒倍数数据切割算法，避免数据浪费
- 动态窗口长度，数据越多分析越准确
- 实时睡眠状态分析和质量评估

使用方法：
- 采用智能数据切割：从数据尾部开始，找到30秒倍数的最大长度
- 支持2小时窗口大小（可配置）
- 最少60秒（2个epoch）即可开始分析

原理：
YASA (Yet Another Spindle Algorithm) 使用SleepStaging类进行自动睡眠分期。
通过智能切割算法，确保每次分析的数据都是30秒的倍数，
充分利用YASA的底层逻辑，实现准确的实时睡眠分析。
数据越长，分析结果越准确，新结果会覆盖旧结果。
"""

import numpy as np
import pandas as pd
from typing import Dict, List, Optional, Tuple, Any
from datetime import datetime, timedelta
import warnings
from collections import deque
import asyncio
import threading
from concurrent.futures import ThreadPoolExecutor

# 尝试导入YASA和MNE库
try:
    import yasa
    import mne
    YASA_AVAILABLE = True
except ImportError:
    YASA_AVAILABLE = False
    yasa = None
    mne = None

# 导入项目模块
from app.models.brain_wave_data import BrainWaveData, TimeSeriesData, FrequencyData
from app.exceptions.device_exceptions import DeviceDataException
from app.config.settings import settings

# 配置日志
import logging
logger = logging.getLogger(__name__)


class YASAAnalyzer:
    """YASA睡眠分析器 - 优化版"""
    
    def __init__(self, 
                 sampling_rate: Optional[int] = None,
                 channel_names: Optional[List[str]] = None):
        """
        初始化YASA分析器
        
        Args:
            sampling_rate: 采样率（Hz），默认512Hz适配NeuroSky设备
            channel_names: 通道名称列表
        """
        # 从配置文件读取参数
        self.window_duration = settings.YASA_WINDOW_DURATION  # 2小时窗口
        self.epoch_duration = settings.YASA_EPOCH_DURATION    # 30秒epoch
        self.min_epochs = settings.YASA_MIN_EPOCHS            # 最少2个epoch
        self.analysis_interval = settings.YASA_ANALYSIS_INTERVAL  # 分析间隔
        
        # 默认使用512Hz采样率适配NeuroSky设备
        self.sampling_rate = sampling_rate or 512
        self.channel_names = channel_names or ['EEG']
        self.yasa_available = YASA_AVAILABLE
        
        # 计算关键参数
        self.epoch_samples = int(self.epoch_duration * self.sampling_rate)  # 30秒的样本数
        self.min_samples = self.min_epochs * self.epoch_samples  # 最少样本数（60秒）
        self.max_samples = int(self.window_duration * self.sampling_rate)  # 最大样本数（2小时）
        
        # 数据存储 - 使用list而不是deque，支持动态长度
        self.data_buffer = []
        self.timestamps_buffer = []
        
        # 分析结果存储 - 按时间点存储多个结果
        self.analysis_results = {}  # {timestamp: [result1, result2, ...]}
        self.last_analysis_time = None
        
        # 异步分析相关
        self.analysis_executor = ThreadPoolExecutor(max_workers=1)
        self.analysis_lock = threading.Lock()
        
        # YASA SleepStaging实例
        self.sleep_staging = None
        
        # 设置MNE日志级别，避免过多的滤波器警告
        if self.yasa_available:
            mne.set_log_level('ERROR')  # 只显示错误信息，隐藏滤波器警告

    def _real_yasa_sleep_staging(self, data: np.ndarray, timestamps: List[datetime], 
                                 use_mne_filtered_data: bool = True) -> Dict[str, Any]:
        """
        使用YASA SleepStaging类进行真实的睡眠分期分析
        
        Args:
            data: 窗口数据（MNE过滤后的数据）
            timestamps: 时间戳列表
            use_mne_filtered_data: 是否使用MNE过滤后的数据（默认True）
            
        Returns:
            Dict: 分析结果
        """
        try:
            logger.info(f"[YASA分析] 开始YASA睡眠分期分析，数据形状: {data.shape}")
            
            # 准备数据格式 - YASA需要 (n_channels, n_samples) 格式
            if data.ndim == 1:
                data = data.reshape(1, -1)
            
            # 创建MNE Info对象
            info = mne.create_info(
                ch_names=self.channel_names,
                sfreq=self.sampling_rate,
                ch_types=['eeg'] * len(self.channel_names)
            )
            
            # 直接使用MNE过滤后的数据创建Raw对象（保持微伏单位）
            raw = mne.io.RawArray(data, info, verbose=False)
            logger.info(f"[YASA分析] MNE Raw对象创建成功，数据长度: {raw.times[-1]:.2f}秒")
            
            # 抑制YASA警告信息
            with warnings.catch_warnings():
                warnings.filterwarnings("ignore")
                
                # 创建YASA SleepStaging实例
                eeg_channel = self.channel_names[0] if self.channel_names else 'EEG'
                sleep_staging = yasa.SleepStaging(raw, eeg_name=eeg_channel)
                
                # 执行睡眠分期预测
                hypnogram = sleep_staging.predict()
                
                # 处理YASA返回结果
                if hasattr(hypnogram, 'hypno'):
                    y_pred = hypnogram.hypno.values if hasattr(hypnogram.hypno, 'values') else hypnogram.hypno
                else:
                    y_pred = hypnogram
                
                logger.info(f"[YASA分析] 预测完成，得到 {len(y_pred)} 个epoch结果")
            
            # 处理epoch结果
            epoch_results = []
            epoch_count = min(len(y_pred), len(data[0]) // self.epoch_samples)
            
            # YASA睡眠阶段映射 - 支持数字和字符串两种格式
            stage_mapping_numeric = {0: 'Wake', 1: 'N1', 2: 'N2', 3: 'N3', 4: 'REM'}
            stage_mapping_string = {'W': 'Wake', 'N1': 'N1', 'N2': 'N2', 'N3': 'N3', 'R': 'REM'}
            
            for i in range(epoch_count):
                epoch_start_index = i * self.epoch_samples
                epoch_end_index = min((i + 1) * self.epoch_samples - 1, len(timestamps) - 1)
                
                if epoch_start_index >= len(timestamps):
                    break
                    
                epoch_start_time = timestamps[epoch_start_index]
                epoch_end_time = timestamps[epoch_end_index]
                
                # 处理不同格式的预测结果
                raw_stage = y_pred[i]
                if isinstance(raw_stage, str):
                    predicted_stage = stage_mapping_string.get(raw_stage, 'Wake')
                else:
                    try:
                        predicted_stage = stage_mapping_numeric.get(int(raw_stage), 'Wake')
                    except (ValueError, TypeError):
                        predicted_stage = 'Wake'
                
                # 获取置信度
                try:
                    proba_result = sleep_staging.predict_proba()
                    if isinstance(proba_result, pd.DataFrame) and i < len(proba_result):
                        confidence = float(proba_result.iloc[i].max())
                    else:
                        confidence = 0.8
                except Exception:
                    confidence = 0.8
                
                epoch_result = {
                    'epoch_index': i,
                    'start_time': epoch_start_time,
                    'end_time': epoch_end_time,
                    'sleep_stage': predicted_stage,
                    'confidence': confidence
                }
                epoch_results.append(epoch_result)
            
            # 计算睡眠质量指标
            sleep_quality = self._calculate_sleep_quality_from_stages(epoch_results)
            
            # 构建分析结果
            analysis_result = {
                'timestamp': timestamps[-1],
                'window_start': timestamps[0],
                'window_end': timestamps[-1],
                'window_duration_seconds': len(timestamps) / self.sampling_rate,
                'total_epochs': epoch_count,
                'epoch_results': epoch_results,
                'current_sleep_stage': epoch_results[-1]['sleep_stage'] if epoch_results else 'Wake',
                'current_confidence': epoch_results[-1]['confidence'] if epoch_results else 0.0,
                'sleep_quality_score': sleep_quality['overall_score'],
                'sleep_efficiency': sleep_quality['efficiency'],
                'deep_sleep_ratio': sleep_quality['deep_sleep_ratio'],
                'rem_sleep_ratio': sleep_quality['rem_sleep_ratio'],
                'analysis_method': 'YASA_SleepStaging',
                'sampling_rate': self.sampling_rate,
                'channels_analyzed': len(self.channel_names),
                'data_length_used': len(data[0])
            }
            
            logger.info(f"[YASA分析] 分析完成，当前睡眠阶段: {analysis_result['current_sleep_stage']}")
            self._store_analysis_result(analysis_result)
            return analysis_result
            
        except Exception as e:
            logger.error(f"[YASA分析] 分析失败: {str(e)}")
            return self._create_default_result(timestamps, f"YASA SleepStaging error: {str(e)}")

    def _calculate_sleep_quality_from_stages(self, epoch_results: List[Dict]) -> Dict[str, float]:
        """
        基于睡眠分期结果计算睡眠质量指标
        
        Args:
            epoch_results: epoch分析结果列表
            
        Returns:
            Dict: 睡眠质量指标
        """
        if not epoch_results:
            return {'overall_score': 0.0, 'efficiency': 0.0, 'deep_sleep_ratio': 0.0, 'rem_sleep_ratio': 0.0}
        
        total_epochs = len(epoch_results)
        stage_counts = {'Wake': 0, 'N1': 0, 'N2': 0, 'N3': 0, 'REM': 0}
        
        # 统计各阶段epoch数量
        for epoch in epoch_results:
            stage = epoch['sleep_stage']
            if stage in stage_counts:
                stage_counts[stage] += 1
        
        # 计算比例
        sleep_epochs = total_epochs - stage_counts['Wake']
        sleep_efficiency = sleep_epochs / total_epochs if total_epochs > 0 else 0.0
        
        deep_sleep_ratio = stage_counts['N3'] / total_epochs if total_epochs > 0 else 0.0
        rem_sleep_ratio = stage_counts['REM'] / total_epochs if total_epochs > 0 else 0.0
        
        # 计算综合睡眠质量分数
        # 深睡眠和REM睡眠比例越高，质量越好
        overall_score = (
            sleep_efficiency * 0.4 +  # 睡眠效率权重40%
            deep_sleep_ratio * 0.35 +  # 深睡眠权重35%
            rem_sleep_ratio * 0.25     # REM睡眠权重25%
        )
        
        return {
            'overall_score': min(overall_score, 1.0),
            'efficiency': sleep_efficiency,
            'deep_sleep_ratio': deep_sleep_ratio,
            'rem_sleep_ratio': rem_sleep_ratio
        }

    def _store_analysis_result(self, result: Dict[str, Any]):
        """
        存储分析结果到历史记录
        
        Args:
            result: 分析结果
        """
        timestamp = result['timestamp']
        
        # 为每个epoch的时间点存储结果
        for epoch in result.get('epoch_results', []):
            epoch_time = epoch['end_time']
            
            if epoch_time not in self.analysis_results:
                self.analysis_results[epoch_time] = []
            
            # 添加新结果
            epoch_result = {
                'sleep_stage': epoch['sleep_stage'],
                'confidence': epoch['confidence'],
                'data_length': result['data_length_used'],
                'analysis_time': timestamp,
                'window_start': result['window_start']
            }
            
            self.analysis_results[epoch_time].append(epoch_result)
            
            # 保持每个时间点最多5个结果
            if len(self.analysis_results[epoch_time]) > 5:
                self.analysis_results[epoch_time] = self.analysis_results[epoch_time][-5:]

    def _create_default_result(self, timestamps: List[datetime], error_msg: str = "") -> Dict[str, Any]:
        """
        创建默认分析结果（分析失败时使用）
        
        Args:
            timestamps: 时间戳列表
            error_msg: 错误信息
            
        Returns:
            Dict: 默认分析结果
        """
        # 检查时间戳数组是否为空
        if not timestamps or len(timestamps) == 0:
            logger.warning("[YASA分析] 时间戳数组为空，使用当前时间作为默认值")
            current_time = datetime.now()
            timestamp_value = current_time
            window_start = current_time
            window_end = current_time
            window_duration = 0.0
            data_length = 0
        else:
            timestamp_value = timestamps[-1]
            window_start = timestamps[0]
            window_end = timestamps[-1]
            window_duration = len(timestamps) / self.sampling_rate
            data_length = len(timestamps)
        
        return {
            'timestamp': timestamp_value,
            'window_start': window_start,
            'window_end': window_end,
            'window_duration_seconds': window_duration,
            'total_epochs': 0,
            
            # 睡眠分期结果
            'epoch_results': [],
            'current_sleep_stage': 'Unknown',
            'current_confidence': 0.0,
            
            # 睡眠质量评估
            'sleep_quality_score': 0.5,
            'sleep_efficiency': 0.5,
            'deep_sleep_ratio': 0.0,
            'rem_sleep_ratio': 0.0,
            
            # 技术指标
            'analysis_method': 'Default',
            'sampling_rate': self.sampling_rate,
            'channels_analyzed': len(self.channel_names),
            'data_length_used': data_length,
            'error_message': error_msg
        }

    def analyze_sleep_stages(self, brain_wave_data_list: List[BrainWaveData]) -> Optional[Dict[str, Any]]:
        """
        分析睡眠阶段 - 兼容tablet_service调用的方法
        
        Args:
            brain_wave_data_list: BrainWaveData对象列表
            
        Returns:
            Optional[Dict]: 睡眠分析结果
        """
        try:
            logger.info(f"=== YASA分析开始 ===")
            logger.info(f"输入数据列表长度: {len(brain_wave_data_list) if brain_wave_data_list else 0}")
            
            if not brain_wave_data_list:
                logger.warning("输入数据列表为空")
                return None
            
            # 提取时域数据和时间戳
            all_data = []
            all_timestamps = []
            
            logger.info("=== 1. 数据提取阶段 ===")
            for i, brain_wave_data in enumerate(brain_wave_data_list):
                logger.info(f"处理第{i+1}个BrainWaveData对象:")
                logger.info(f"  - 类型: {type(brain_wave_data)}")
                logger.info(f"  - 有time_series: {hasattr(brain_wave_data, 'time_series') and brain_wave_data.time_series is not None}")
                
                if brain_wave_data.time_series and brain_wave_data.time_series.data is not None:
                    # 正确处理时域数据：每个BrainWaveData包含一个时间序列
                    data = brain_wave_data.time_series.data
                    logger.info(f"  - 数据类型: {type(data)}")
                    logger.info(f"  - 数据形状: {data.shape if hasattr(data, 'shape') else 'N/A'}")
                    logger.info(f"  - 数据长度: {len(data) if hasattr(data, '__len__') else 'N/A'}")
                    
                    if isinstance(data, np.ndarray):
                        if data.ndim == 0:  # 标量 - 单个采样点
                            all_data.append(float(data))
                            logger.info(f"  - 添加单个采样点: {float(data)}")
                        elif data.ndim == 1:  # 一维数组 - 时间序列
                            # 这是正确的时域数据：一个完整的时间序列
                            all_data.extend(data.flatten())
                            logger.info(f"  - 添加时间序列数据，长度: {len(data)}")
                        elif data.ndim == 2:  # 二维数组 - 多通道时间序列
                            # 假设第一个维度是通道，第二个是时间
                            if data.shape[0] == 1:  # 单通道
                                all_data.extend(data[0])
                                logger.info(f"  - 添加单通道时间序列，长度: {data.shape[1]}")
                            else:  # 多通道，取第一个通道
                                all_data.extend(data[0])
                                logger.info(f"  - 添加多通道时间序列第一通道，长度: {data.shape[1]}")
                        else:
                            # 高维数组，展平处理
                            all_data.extend(data.flatten())
                            logger.info(f"  - 添加高维数组（展平），长度: {len(data.flatten())}")
                    else:
                        # 非numpy数组，尝试转换
                        try:
                            if hasattr(data, '__iter__') and not isinstance(data, str):
                                # 可迭代对象（如列表）
                                all_data.extend([float(x) for x in data])
                                logger.info(f"  - 添加可迭代数据，长度: {len(data)}")
                            else:
                                # 单个数值
                                all_data.append(float(data))
                                logger.info(f"  - 添加单个数值: {float(data)}")
                        except (ValueError, TypeError) as e:
                            logger.warning(f"  - 数据转换失败: {e}")
                            continue
                    
                    # 修复：为每个数据点生成对应的时间戳
                    # 根据实际添加到all_data中的数据点数量来生成时间戳
                    base_timestamp = brain_wave_data.time_series.timestamp
                    data_points_added = 0
                    
                    # 计算实际添加的数据点数量
                    if isinstance(data, np.ndarray):
                        if data.ndim == 1:
                            data_points_added = len(data)
                        elif data.ndim == 2:
                            data_points_added = data.shape[1] if data.shape[0] == 1 else data.shape[1]
                        else:
                            data_points_added = len(data.flatten())
                    else:
                        if hasattr(data, '__iter__') and not isinstance(data, str):
                            data_points_added = len(data)
                        else:
                            data_points_added = 1
                    
                    # 为每个数据点生成时间戳
                    if data_points_added > 1:
                        # 多个数据点，生成时间戳序列
                        for j in range(data_points_added):
                            # 每个采样点间隔 1/512 秒
                            sample_timestamp = base_timestamp + timedelta(seconds=j/self.sampling_rate)
                            all_timestamps.append(sample_timestamp)
                        logger.info(f"  - 生成时间戳序列，数量: {data_points_added}")
                    else:
                        # 单个数据点，使用原始时间戳
                        all_timestamps.append(base_timestamp)
                        logger.info(f"  - 使用原始时间戳")
                else:
                    logger.warning(f"  - 第{i+1}个对象没有有效的time_series数据")
            
            logger.info(f"=== 数据提取完成 ===")
            logger.info(f"总数据点数: {len(all_data)}")
            logger.info(f"时间戳数量: {len(all_timestamps)}")
            
            # 检查数据量是否足够进行YASA分析
            # 根据YASA标准，需要至少30秒的连续时域数据（1个epoch）
            min_required_samples = self.epoch_samples  # 30秒 * 512Hz = 15360个样本
            
            if len(all_data) < min_required_samples:
                logger.warning(f"数据不足进行YASA分析")
                logger.warning(f"当前样本数: {len(all_data)}")
                logger.warning(f"最小需要: {min_required_samples} (30秒 @ {self.sampling_rate}Hz)")
                logger.warning(f"建议: 需要至少30秒的连续EEG数据进行睡眠分期")
                
                # 对于数据不足的情况，生成模拟结果
                logger.info("数据不足，生成基于频域分析的模拟结果")
                return self._generate_fallback_result(all_data, all_timestamps)
            
            # 转换为numpy数组
            data_array = np.array(all_data)
            logger.info(f"=== 2. 数据切块阶段 ===")
            logger.info(f"数据数组形状: {data_array.shape}")
            logger.info(f"数据数组类型: {data_array.dtype}")
            logger.info(f"数据范围: [{np.min(data_array):.2f}, {np.max(data_array):.2f}]")
            
            # 修复：直接使用完整数据进行分析，不通过缓冲区
            logger.info(f"=== 3. YASA分析执行阶段 ===")
            logger.info("直接使用完整数据进行YASA分析...")
            
            # 检查数据量是否足够进行YASA分析
            total_samples = len(data_array)
            logger.info(f"直接分析 - 总样本数: {total_samples}")
            logger.info(f"直接分析 - 最小样本数要求: {self.min_samples}")
            logger.info(f"直接分析 - 每个epoch样本数: {self.epoch_samples}")
            
            if total_samples < self.min_samples:
                logger.warning(f"直接分析 - 数据不足，总样本数 {total_samples} < 最小要求 {self.min_samples}")
                return self._generate_fallback_result(all_data, all_timestamps)
            
            # 计算最大可用的完整epoch数量
            max_epochs = total_samples // self.epoch_samples
            logger.info(f"直接分析 - 可用的最大epoch数: {max_epochs}")
            
            # 确保至少有最少epoch数量
            if max_epochs < self.min_epochs:
                logger.warning(f"直接分析 - epoch数不足，可用 {max_epochs} < 最小要求 {self.min_epochs}")
                return self._generate_fallback_result(all_data, all_timestamps)
            
            # 计算实际使用的样本数（30秒的倍数）
            used_samples = max_epochs * self.epoch_samples
            logger.info(f"直接分析 - 实际使用的样本数: {used_samples}")
            
            # 从尾部开始切割（使用最新的数据）
            start_index = total_samples - used_samples
            logger.info(f"直接分析 - 切割起始索引: {start_index}")
            
            cut_data = data_array[start_index:]
            cut_timestamps = all_timestamps[start_index:]
            
            logger.info(f"直接分析 - 切割后数据形状: {cut_data.shape}")
            logger.info(f"直接分析 - 切割后时间戳数量: {len(cut_timestamps)}")
            
            # 关键修复：调用YASA分析方法，明确指定使用MNE过滤后的数据
            result = self._real_yasa_sleep_staging(cut_data, cut_timestamps, use_mne_filtered_data=True)
            logger.info(f"直接YASA分析返回结果: {result}")
            
            if result:
                logger.info("=== YASA分析成功 ===")
                return result
            else:
                logger.warning("YASA分析未返回结果")
                return None
                
        except Exception as e:
            logger.error(f"analyze_sleep_stages异常: {str(e)}", exc_info=True)
            logger.info(f"=== YASA分析异常结束 ===")
            return None

    def _generate_fallback_result(self, data: List[float], timestamps: List[datetime]) -> Dict[str, Any]:
        """
        当数据不足进行完整YASA分析时，生成基于简单频域分析的后备结果
        
        Args:
            data: 数据点列表
            timestamps: 时间戳列表
            
        Returns:
            Dict: 简化的分析结果
        """
        try:
            if len(data) < 2:
                return self._create_default_result(timestamps, "数据点不足")
            
            # 转换为numpy数组
            data_array = np.array(data)
            
            # 简单的频域分析
            from scipy import signal
            
            # 计算功率谱密度
            freqs, psd = signal.welch(data_array, fs=self.sampling_rate, nperseg=min(len(data_array), 256))
            
            # 定义频段
            delta_band = (0.5, 4)
            theta_band = (4, 8)
            alpha_band = (8, 13)
            beta_band = (13, 30)
            
            # 计算各频段功率
            def get_band_power(freqs, psd, band):
                idx = np.logical_and(freqs >= band[0], freqs <= band[1])
                return np.trapz(psd[idx], freqs[idx]) if np.any(idx) else 0.0
            
            delta_power = get_band_power(freqs, psd, delta_band)
            theta_power = get_band_power(freqs, psd, theta_band)
            alpha_power = get_band_power(freqs, psd, alpha_band)
            beta_power = get_band_power(freqs, psd, beta_band)
            
            total_power = delta_power + theta_power + alpha_power + beta_power
            
            if total_power > 0:
                # 基于频段功率比例判断睡眠状态
                alpha_ratio = alpha_power / total_power
                beta_ratio = beta_power / total_power
                
                if alpha_ratio > 0.3 or beta_ratio > 0.4:
                    sleep_stage = "Wake"
                    confidence = 0.6
                elif theta_power / total_power > 0.3:
                    sleep_stage = "N1"
                    confidence = 0.5
                else:
                    sleep_stage = "N2"
                    confidence = 0.4
            else:
                sleep_stage = "Wake"
                confidence = 0.3
            
            # 构建结果
            result = {
                "sleep_stage": sleep_stage,
                "confidence": confidence,
                "frequency_bands": {
                    "Delta": delta_power / total_power if total_power > 0 else 0,
                    "Theta": theta_power / total_power if total_power > 0 else 0,
                    "Alpha": alpha_power / total_power if total_power > 0 else 0,
                    "Beta": beta_power / total_power if total_power > 0 else 0,
                },
                "data_points": len(data),
                "analysis_duration": "N/A",
                "timestamp": timestamps[-1] if timestamps else datetime.now(),
                "method": "fallback_frequency_analysis",
                "note": "数据不足，使用简化频域分析"
            }
            
            logger.info(f"生成后备分析结果: {sleep_stage}, 置信度: {confidence:.3f}")
            return result
            
        except Exception as e:
            logger.error(f"生成后备结果时出错: {str(e)}")
            return self._create_default_result(timestamps, f"后备分析失败: {str(e)}")